use std::{
    collections::HashMap,
    fmt,
    sync::{Arc, RwLock},
    task::{Context, Poll},
};

use axum::{
    extract::{Path, State},
    routing::{delete, get, Route},
    Router,
};
use tower::{service_fn, BoxError, Service, ServiceBuilder};
type SharedState = Arc<RwLock<AppState>>;
#[derive(Default)]
struct AppState {
    db: HashMap<String, String>,
}

async fn list_keys(State(state): State<SharedState>) -> String {
    let db = &state.read().unwrap().db;

    db.keys()
        .map(|key| key.to_string())
        .collect::<Vec<String>>()
        .join("\n")
}

fn admin_routes() -> Router<SharedState> {
    async fn delete_all_keys(State(state): State<SharedState>) {
        state.write().unwrap().db.clear();
    }

    async fn remove_key(Path(key): Path<String>, State(state): State<SharedState>) {
        state.write().unwrap().db.remove(&key);
    }

    Router::new()
        .route("/keys", delete(delete_all_keys))
        .route("/key/:key", delete(remove_key))
    // Require bearer auth for all admin routes
}

#[derive(Debug, Clone)]
pub struct LogService<S> {
    service: S,
}

impl<S, Request> Service<Request> for LogService<S>
where
    S: Service<Request>,
    Request: fmt::Debug,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // Insert log statement here or other functionality
        println!("befor request = {:?}", request);
        let r = self.service.call(request);
        return r;
    }
}

#[derive(Debug, Clone)]
pub struct LogService1<S> {
    service: S,
}

impl<S, Request> Service<Request> for LogService1<S>
where
    S: Service<Request>,
    Request: fmt::Debug,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(cx)
    }

    fn call(&mut self, request: Request) -> Self::Future {
        // Insert log statement here or other functionality
        println!("befor request = {:?}", request);
        let r = self.service.call(request);
        return r;
    }
}

#[tokio::main]
async fn main() {
    let shared_state = SharedState::default();

    let app = Router::new()
        .route("/keys", get(list_keys))
        .nest("/admin", admin_routes())
        .layer(
            ServiceBuilder::new()
                .layer_fn(|service| LogService { service })
                .layer_fn(|service| LogService1 { service }),
        )
        .with_state(Arc::clone(&shared_state));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000")
        .await
        .unwrap();
    axum::serve(listener, app).await.unwrap();
}

#[cfg(test)]
mod tests {
    use std::{
        borrow::Borrow,
        fmt,
        task::{Context, Poll},
    };

    use tower::{service_fn, BoxError, Layer, Service, ServiceBuilder, ServiceExt};

    pub struct LogLayer {
        target: &'static str,
    }

    impl<S> Layer<S> for LogLayer {
        type Service = LogService<S>;

        fn layer(&self, service: S) -> Self::Service {
            LogService {
                target: self.target,
                service,
            }
        }
    }

    // This service implements the Log behavior
    pub struct LogService<S> {
        target: &'static str,
        service: S,
    }

    impl<S> Service<i32> for LogService<S>
    where
        S: Service<&'static str>,
        // Request: fmt::Debug,
    {
        type Response = S::Response;
        type Error = S::Error;
        type Future = S::Future;

        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.service.poll_ready(cx)
        }

        fn call(&mut self, request: i32) -> Self::Future {
            // Insert log statement here or other functionality
            println!("befor request = {}, target = {:?}", request, self.target);
            let r = self.service.call(self.target);
            println!("after  target = {:?}", self.target);
            return r;
        }
    }

    #[tokio::test]
    // #[test]
    pub async fn test2() {
        async fn handle(request: &'static str) -> Result<&'static str, BoxError> {
            println!("11111111111:{}", request);
            Ok(request)
        }

        let mut f = ServiceBuilder::new()
            .layer_fn(|service| LogService {
                target: "test",
                service,
            })
            .service_fn(handle);

        let f = f.oneshot(1).await;
        println!("f:{}", f.unwrap());
    }
}
